from metagpt.team import Team
from test.utils import stream_jsonl, run_model, RESULT_DIR
import asyncio
import math
import os
import threading
import datetime
import shutil
import re
import subprocess
import time

# 最多线程数量，防止使用 API 并发数过高
MAX_THREAD_NUM = 5
pass_at_1 = 0.0
resume = 0
total_task_num = 0
start_time = 0
lock = threading.Lock()
lock2 = threading.Lock()
lock3 = threading.Lock()

# result_dir is like "./results/humanevalx_pass@1_15:00:00/"   
def run_test(task, n: int, k: int, result_dir: str, use_naive=False, use_rag=True):
    if "c_code" in task.keys() and "rust_test" in task.keys():
        test_ordinary(task, n, k, result_dir, use_naive, use_rag)
    elif "client_code" in task.keys() and "server_code" in task.keys():
        test_client_server(task, n, k, result_dir, use_naive, use_rag)
    else:
        print("Invalid task format")
        
    
def test_ordinary(task, n: int, k: int, result_dir: str, use_naive=False, use_rag=True):
    global pass_at_1
    global total_task_num
    global start_time
    pass_time = 0
    task_id = str(task["id"])
    base_dir = result_dir + str(task_id) + "/"
    os.makedirs(base_dir)
    cpp_file_name = base_dir + "main_" + task_id + ".cpp"
    with open(cpp_file_name, "w") as f:
        f.write(task["c_code"])
        
    with open("unittest.rs", "w") as f:
        f.write(task["rust_test"])
        
    with open(base_dir + task_id + ".log", "w") as log:
        for i in range(n):
            # lock3.acquire(blocking=True)
            (code_text, compiler_message) =  asyncio.run(run_model(
                c_code=task["c_code"],
                use_naive=use_naive,
                use_rag=use_rag,
                investment=5.0,
                max_round=20))
            if os.path.exists("epoch.txt"):
                with open("epoch.txt", "r") as f1:
                    epoch = f1.read()
                with open(base_dir + "server_" + str(i) + "epoch.txt", "w") as f1:
                    f1.write(epoch)
            with open("unittest.rs", "r") as f:
                unittest = f.read()
            # lock3.release()
            file_name = base_dir + "main" + task_id + str(i) + "time" + ".rs"
            with open(file_name, "w") as f:
                f.write(code_text + unittest)
            pattern = r"error"
            match = re.search(pattern, compiler_message, re.DOTALL)
            if not match:
                # 通过编译
                log.write("#" + str(i) + " Success\n")
            else:
                log.write("#" + str(i) + " Compile failed\n")
        # lock.acquire(blocking=True)
        value = 1.0 - (math.comb(n - pass_time, k) / math.comb(n, k))
        pass_at_1 *= total_task_num
        pass_at_1 += value
        total_task_num += 1
        pass_at_1 /= total_task_num
        log.write(f"Run time: {round(time.time() - start_time, 3)} s\n")
        log.write(f"pass@1: {pass_at_1}")
        # lock.release()
        
def test_client_server(task, n: int, k: int, result_dir: str, use_naive=False, use_rag=True):
    global pass_at_1
    global total_task_num
    global start_time
    pass_time = 0
    task_id = str(task["id"])
    base_dir = result_dir + str(task_id) + "/"
    os.makedirs(base_dir)
    client_file_name = base_dir + "client_" + task_id + ".cpp"
    with open(client_file_name, "w") as f:
        f.write(task["client_code"])
    server_file_name = base_dir + "server_" + task_id + ".cpp"
    with open(server_file_name, "w") as f:
        f.write(task["server_code"])
    
    if os.path.exists("unittest.rs"):
        os.remove("unittest.rs")
        
    with open(base_dir + task_id + ".log", "w") as log:
        for i in range(n):
            # lock3.acquire(blocking=True)
            (client_code_text, client_compiler_message) =  asyncio.run(run_model(
                c_code=task["client_code"],
                use_naive=use_naive,
                use_rag=use_rag,
                investment=5.0,
                max_round=20))
            if os.path.exists("epoch.txt"):
                with open("epoch.txt", "r") as f1:
                    epoch = f1.read()
                with open(base_dir + "server_" + str(i) + "epoch.txt", "w") as f1:
                    f1.write(epoch)
            (server_code_text, server_compiler_message) =  asyncio.run(run_model(
                c_code=task["server_code"],
                use_naive=use_naive,
                use_rag=use_rag,
                investment=5.0,
                max_round=20))
            if os.path.exists("epoch.txt"):
                with open("epoch.txt", "r") as f1:
                    epoch = f1.read()
                with open(base_dir + "server_" + str(i) + "epoch.txt", "w") as f1:
                    f1.write(epoch)
            # lock3.release()
            file_name = base_dir + "client_main" + task_id + str(i) + "time" + ".rs"
            with open(file_name, "w") as f:
                f.write(client_code_text)
            file_name = base_dir + "server_main" + task_id + str(i) + "time" + ".rs"
            with open(file_name, "w") as f:
                f.write(server_code_text)
            pattern = r"error"
            match1 = re.search(pattern, client_compiler_message, re.DOTALL)
            match2 = re.search(pattern, server_compiler_message, re.DOTALL)
            if not match1 and not match2:
                # 通过编译
                log.write("#" + str(i) + " Success\n")
            else:
                log.write("#" + str(i) + " Compile failed\n")
        # lock.acquire(blocking=True)
        value = 1.0 - (math.comb(n - pass_time, k) / math.comb(n, k))
        pass_at_1 *= total_task_num
        pass_at_1 += value
        total_task_num += 1
        pass_at_1 /= total_task_num
        log.write(f"Run time: {round(time.time() - start_time, 3)} s\n")
        log.write(f"pass@1: {pass_at_1}")
        # lock.release()
        
    
def output_only_code(k: int = 1, use_naive=False, use_rag=True):
    global start_time
    global resume
    n = 5
    
    # ths = []
    
    now = (datetime.datetime.now() + datetime.timedelta(hours=8)).strftime('%H:%M:%S')
    result_dir = RESULT_DIR + "humanevalx_pass@" + str(k) + "_" + now + "/"
    if os.path.exists(result_dir):
        shutil.rmtree(result_dir)
    os.makedirs(result_dir)
    start_time = time.time()
    j = 0
    for task in stream_jsonl("./testcases/cpptrait/cpptrait.jsonl.gz"):
        if j < resume:
            j += 1
            continue
        run_test(task, n, k, result_dir, use_naive, use_rag)
